#!/usr/bin/python
# Author: Jimmy Saw
# Date: 10-16-2011
# This program parses 454Scaffolds.txt and mummer coordinate files.
# Needs Newbler version 2.6 run with options to produce "Output trimmed reads" and "Output scaffolds"
# Then run nucmer of filled scaffold contigs "454ScaffoldContigs.fna" against mates.fasta file
# show-coords -r -c -l -T prefix.delta > prefix.coords
# Usage: python newblerFilledScaffolds.py 454Scaffolds.txt newbler_scaffolds_vs_mates.coords

import sys
import re
import itertools
import matplotlib.pyplot as plt
import pylab

#regex
ctg = re.compile('sctg_\d+_\d+')
left = re.compile('(.*)_left') #Edit for your mate pair names
right = re.compile('(.*)_right') #Edit for your mate pair names
odd = re.compile('.*[1|3|5|7|9]$')
even = re.compile('.*[2|4|6|8|0]$')

scaffile = open(sys.argv[1], "rU")
sfl = scaffile.readlines()

pairfile = open(sys.argv[2], "rU")
pfl = pairfile.readlines()

#Work on 454Scaffolds.txt file
scaflist = []
ctglist = []
gaplist = []

x1 = []
x2 = []
y = []
ctgx = []
ctgy = []

#parse scaffold file and put coordinates into place
for index, line in enumerate(sfl):
    i = line.split('\t')
    scaf = i[0]
    scafstart = int(i[1])
    scafstop = int(i[2])
    ctgname = i[5]
    ctgstart = i[6]
    ctgstop = i[7].rstrip()
    if ctg.match(ctgname):
        ctgorient = i[8].rstrip()
        ctglen = int(ctgstop) - int(ctgstart) + 1
        ctglist.append((ctgname, scafstart, scafstop, int(ctgstart), int(ctgstop), ctgorient, ctglen))
        x1.append(scafstart)
        x2.append(scafstop)
        if odd.match(str(ctgname)):
            y.append(4)
        else:
            y.append(4.2)
    else:
        gaplength = int(ctgname)
        gaplist.append((scafstart, scafstop, gaplength))

scaflength = ctglist[-1][2]
scafx = [1, scaflength]
scafy = [3, 3]

a = [x1, x2]
b = [y, y]

#Work on MUMMER coordinate file
templatelist = []
linklist = []
contigslist = []
leftlist = []
rightlist = []
matepairs = []
placementlist = {}
#placementdict = {}

for index, line in enumerate(pfl[4:]):
    i = line.split('\t')
    readname = i[-1].rstrip()
    ctgname = i[-2]
    if left.match(readname):
        template = left.match(readname)
        templatelist.append(template.group(1))
    elif right.match(readname):
        template = right.match(readname)
        templatelist.append(template.group(1))
    contigslist.append(ctgname)
    placementlist[readname] = ctgname 
    #Eg: placement['GM6SIKE01A8JWF_right'] = 'sctg_0001_0001'

uniquectglist = list(set(contigslist))
uniquetemplist = list(set(templatelist))

gspairs = []
ctgpairs = []
for i in uniquetemplist:
    lread = i + "_left"
    rread = i + "_right"
    matepairs.append((lread, rread))
    if lread in placementlist and rread in placementlist:
        if placementlist[lread] != placementlist[rread]:
            gspairs.append((lread, rread)) #gap-spanning read pairs
            ctgpairs.append((placementlist[lread], placementlist[rread])) #linked contigs

contigpairs = list(itertools.permutations(uniquectglist, 2))

pairdict = {}

for i in contigpairs:
    pairdict[i[0] + "_" + i[1]] = 0
    #pairdict[(i[0], i[1])] = 0
    #Eg: pairdict['sctg_0001_0001_sctg_0001_0002'] = 0

for i in ctgpairs:
    pair = i[0] + "_" + i[1]
    if pair in pairdict:
        pairdict[pair] += 1

fig = plt.figure(1, figsize=(14,6))
plt.plot(scafx, scafy, color='#666600', marker='|', markersize=8.0,
    markeredgecolor='black', linestyle='-', linewidth=2.0)
plt.plot(a, b, color='purple', linestyle='-', linewidth=1.5)
plt.plot(x2, y, 'r>', mec='purple')

contiglayout = ""

for i, j in enumerate(gaplist):
    gapmiddle = j[0] + ((j[1] - j[0]) / 2)
    if odd.match(str(i)):
        plt.text(gapmiddle, 4.8, str(j[2]), rotation=45, fontsize=8)
    else:
        plt.text(gapmiddle, 5, str(j[2]), rotation=45, fontsize=8)

for i, j in enumerate(ctglist):
    ctgmiddle = j[1] + ((j[2] - j[1]) / 2)
    gapcenter = 0
    linkcount = 0
    link2print = ""
    if i < len(ctglist)-1:
        currctg = ctglist[i]
        nextctg = ctglist[i+1]
        pair = currctg[0] + "_" + nextctg[0]
        if pair in pairdict:
            linkcount = pairdict[currctg[0] + "_" + nextctg[0]]
            link2print =  currctg[0][-4:] + "_" + nextctg[0][-4:] + ": " + str(linkcount)
            if linkcount != 0:
                contiglayout += currctg[0] + " " + str(linkcount) + " " + nextctg[0] + "\n"
        gapcenter = currctg[2]
        if linkcount == 0:
            newpair = currctg[0] + "_" + ctglist[i+2][0] #check contig two gaps away
            if newpair in pairdict:
                linkcount = pairdict[newpair]
                link2print =  currctg[0][-4:] + "_" + ctglist[i+2][0][-4:] + ": " + str(linkcount)
                if linkcount != 0:
                    contiglayout += currctg[0] + " " + str(linkcount) + " " + ctglist[i+2][0] + "\n"
            pair = newpair
        if linkcount == 0:
            newpair = currctg[0] + "_" + ctglist[i+3][0] #check contig three gaps away
            if newpair in pairdict:
                linkcount = pairdict[newpair]
                link2print =  currctg[0][-4:] + "_" + ctglist[i+3][0][-4:] + ": " + str(linkcount)
                if linkcount != 0:
                    contiglayout += currctg[0] + " " + str(linkcount) + " " + ctglist[i+3][0] + "\n"
            pair = newpair
        if linkcount == 0:
            newpair = currctg[0] + "_" + ctglist[i+4][0] #check contig four gaps away
            if newpair in pairdict:
                linkcount = pairdict[newpair]
                link2print =  currctg[0][-4:] + "_" + ctglist[i+4][0][-4:] + ": " + str(linkcount)
                if linkcount != 0:
                    contiglayout += currctg[0] + " " + str(linkcount) + " " + ctglist[i+4][0] + "\n"
            pair = newpair
    else:
        pass

    #if odd.match(str(i)):
    if odd.match(j[0]):
        plt.text(ctgmiddle, 4, str(j[0][-4:]), rotation=90, fontsize=8, color='#458000')
        plt.text(gapcenter, 4, link2print, rotation=90, fontsize=8, color='red')
        plt.text(ctgmiddle, 3.90, str(j[-1]), fontsize=8, color='purple')
    else:
        plt.text(ctgmiddle, 4.2, str(j[0][-4:]), rotation=90, fontsize=8, color='#458000')
        plt.text(gapcenter, 4.2, link2print, rotation=90, fontsize=8, color='red')
        plt.text(ctgmiddle, 4.1, str(j[-1]), fontsize=8, color='purple')

print contiglayout

plt.text(3000000, 3.5, "Purple = Newbler Contigs", fontsize=10, color='purple')
plt.text(3000000, 3.3, "Red = Count of links", fontsize=10, color='red')
plt.text(3000000, 3.05, "Black = Gap size", fontsize=10, color='black')
#plt.text(3000000, 2.9, "Blue = Contig size", fontsize=10, color='purple')

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)

plt.axis([0, scaflength, 2, 6])
#plt.title('Newbler Filled Contig Scaffolds: sctg_xxxx_xxxx')
plt.xlabel('Scaffold coordinates')

plt.grid(True)

scaffile.close()
pairfile.close()

plt.show()
